import h5py
import torch
import torch.utils.data as data
import random
import utils.utils as utils
from .database import Database
#一共 600 张图片   70%训练:840  30%测试 : 360

def cut_picture(M1,M2,M3,M4, dH=128,dW=128):    #batch==1   输入一个数组
    sizes=M1[0].size()
    outs1=[]
    outs2=[]
    outs3=[]
    outs4=[]
    H=sizes[-2]
    W=sizes[-1]

    x=random.randint(0,H-dH)
    y=random.randint(0,W-dW)
    for i1 in M1:
        outs1.append(i1[...,x:x+dH,y:y+dW])
    for i2 in M2:
        outs2.append(i2[...,x:x+dH,y:y+dW])
    for i3 in M3:
        outs3.append(i3[...,x:x+dH,y:y+dW])
    for i4 in M4:
        outs4.append(i4[...,x:x+dH,y:y+dW])
    return  outs1,outs2,outs3,outs4


class Ballet(Database):
    def __init__(self,cfg,flag='train'):
        self.img_size = cfg.img_size
        self.framnumber=98
        self.height = 768
        self.width = 1024
        self.traincames=[1,2,4,5,6]
        self.testcames=[3]
        self.cames=[1,2,3,4,5,6]
        # if utils.ip=='176':
        #     self.h5path=cfg.data_root
        # if utils.ip=='175':
        #     self.h5path=cfg.data_root
        self.h5path=cfg.data_root
        self.flag=flag
        if flag=='train':
            self.cames1= self.traincames
        if flag=='test':
            self.cames1= self.testcames
    def getdata(self,cames,fram,):
        cams1='cam{}'.format(cames)
        with h5py.File(self.h5path.format(cames),'r') as f:
            mask=torch.from_numpy(f[cams1]['mask'][fram]/255).unsqueeze(0).float()
            pictureL=torch.from_numpy(f[cams1]['imgL'][fram]/255).permute(2,0,1).float()  #名称为image 
            pictureR =torch.from_numpy(f[cams1]['imgR'][fram]/255).permute(2,0,1).float() #名称为image
            dpointL =torch.from_numpy(f[cams1]['Lerror'][fram]).permute(2,0,1).float()  #名称为image
            dpointR=torch.from_numpy(f[cams1]['Rerror'][fram]).permute(2,0,1).float()  #名称为image
            target=torch.from_numpy(f[cams1]['target'][fram]/255).permute(2,0,1).float()   #名称为imag
        return mask,pictureL,pictureR,dpointL,dpointR,target